import os
import xml.etree.ElementTree as ET
from collections import defaultdict

def count_labels_from_xml(xml_folder):
    """
    统计LabelImg生成的XML文件中的标签数量。
    参数:
        xml_folder (str): 包含XML标注文件的文件夹路径。
    返回:
        dict: 每个标签及其对应的数量。
    """
    # 初始化一个字典来存储标签计数
    label_count = defaultdict(int)
    
    # 遍历文件夹中的所有XML文件
    for xml_file in os.listdir(xml_folder):
        if xml_file.endswith('.xml'):
            xml_path = os.path.join(xml_folder, xml_file)
            
            # 解析XML文件
            tree = ET.parse(xml_path)
            root = tree.getroot()
            
            # 遍历所有的<object>节点
            for obj in root.findall('object'):
                label = obj.find('name').text
                label_count[label] += 1
    
    return label_count

def main():
    # 输入XML文件夹路径
    xml_folder = input("请输入包含XML标注文件的文件夹路径: ")
    
    # 检查路径是否存在
    if not os.path.exists(xml_folder):
        print("路径不存在，请检查后重试！")
        return
    
    # 统计标签数量
    label_count = count_labels_from_xml(xml_folder)
    
    # 输出结果
    print("\n标签统计结果：")
    for label, count in label_count.items():
        print(f"{label}: {count}")
    
    # 可选：保存结果到文件
    save_to_file = input("\n是否将结果保存到文件？(y/n): ").strip().lower()
    if save_to_file == 'y':
        output_file = os.path.join(xml_folder, "label_count.txt")
        with open(output_file, 'w') as f:
            for label, count in label_count.items():
                f.write(f"{label}: {count}\n")
        print(f"结果已保存到文件：{output_file}")

if __name__ == "__main__":
    main()